#!/usr/bin/env python
# encoding: utf-8

from core.log import Log
from core.log import color
from core.banner import banner

import sys
import string
import os
import hashlib
import readline
import code
import atexit
import json
import time
import signal
import importlib

def setup():
    history_file = "./.history"
    if not os.path.exists(history_file):
        open(history_file, 'a+').close()

    readline.read_history_file(history_file)
    readline.set_history_length(history_length)
    atexit.register(readline.write_history_file, history_file)

    readline.parse_and_bind('set enable-keypad on')

    readline.set_completer(complete)
    readline.set_completer_delims(' \t\n;')
    readline.parse_and_bind("tab: complete")

def md5(content):
    return hashlib.md5(content).hexdigest()

def show_help():
    print "Usage : "
    print "        python %s" % (sys.argv[0])
    print "Author : "
    print "        WangYihang <wangyihanger@gmail.com>"
    print "Github : "
    print "        https://github.com/wangyihang/exploit-framework"

def core_commands():
    print "Core Commands"
    print "============="
    print "\tCommand\tDescription"
    print "\t-------\t-----------"
    print "\thelp\tshow help"
    print "\tversion\tshow version"
    print "\tuse\tSelects a module by name"
    print "\tshow\tDisplays modules of a given type, or all modules"
    print "\tsearch\tSearches module names and descriptions"
    print "\tback\tMove back from the current context"
    print "\tquit\tquit"
    print ""

def module_command():
    print "Module Commands"
    print "============="
    print "\tCommand\tDescription"
    print "\t-------\t-----------"
    print "\toptions\tDisplays global options or for one or more modules"
    print "\tinfo\tDisplays information about one or more modules"
    print ""

def main_help():
    core_commands()
    module_command()

def signal_handler(ignum, frame):
    print ""
    Log.info("Enter : 'q|quit|exit' to shutdown the program!")

def reset_context():
    return "Framework"

def main():
    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)
    banner()

    LOCAL_COMMAND_FLAG = True

    CONTEXT = reset_context()
    while True:
        command = (raw_input("[%s]=> " % (color.red(CONTEXT))) or "help")
        if command == "h" or command == "help" or command == "?":
            main_help()
        elif command == "version":
            Log.info("Version: 0.0.1")
        elif command == "show":
            print "%s" % (color.purple("------\t\t------"))
            print "%s" % (color.purple("Vendor\t\tModule"))
            print "%s" % (color.purple("------\t\t------"))
            exploit_path = "./exploit/"
            vendors = os.listdir(exploit_path)
            for vendor in vendors:
                full_path = exploit_path + vendor
                if os.path.isdir(full_path):
                    # Log.info("%s" % ("-" * 0x20))
                    # Log.info("Vendor: %s" % (vendor))
                    exploit_files = os.listdir(full_path)
                    number = 0
                    for exploit_file in exploit_files:
                        if exploit_file.endswith(".py") and exploit_file != "__init__.py":
                            # Log.info("%s => exploit.%s.%s" % (exploit_file, vendor, exploit_file.replace(".py", "")))
                            if len(vendor) > 8:
                                print "%s" % (color.cyan("%s\t%s" % (vendor, exploit_file.replace(".py", ""))))
                            else:
                                print "%s" % (color.cyan("%s\t\t%s" % (vendor, exploit_file.replace(".py", ""))))
                            number += 1
                    # Log.info("%d exploits" % (number))
            print "%s" % (color.purple("---------"))
            print "%s" % (color.purple(" Example"))
            print "%s" % (color.purple("---------"))
            print "%s" % (color.cyan("use exploit.%s.%s" % (vendor, exploit_file.replace(".py", ""))))
        elif command.startswith("use "):
            module_name = command.split(" ")[1]
            Log.info("Loading module: %s" % (module_name))
            try:
                module = importlib.import_module(module_name)
            except Exception as e:
                Log.error(str(e))
                continue
            CONTEXT = module_name
            exploit = module.Exploit()
            exploit.show_info()
            Log.info("%s" % ("-" * 0x40))
            exploit.show_options()
            while True:
                module_command = (raw_input("[%s]=> " % (color.red(CONTEXT))) or "help")
                if module_command == "help":
                    main_help()
                    continue
                if module_command.startswith("set "):
                    if len(module_command.split(" ")) == 3:
                        key = module_command.split(" ")[1]
                        value = module_command.split(" ")[2]
                        exploit.set_config(key, value)
                    else:
                        Log.error("Check your input!")
                        Log.info("Example: \n\tset [KEY] [VALUE]")
                elif module_command == "options":
                    exploit.show_options()
                elif module_command == "info":
                    exploit.show_info()
                elif module_command == "exploit":
                    try:
                        exploit.exploit()
                    except Exception as e:
                        Log.error(str(e))
                elif module_command == "quit" or module_command == "q" or module_command == "exit" or module_command == "back":
                    break
                else:
                    main_help()
            CONTEXT = reset_context()
        elif command == "q" or command == "quit" or command == "exit":
            Log.info("Quiting...")
            break
        else:
            Log.error("Unsupported function!")
            if LOCAL_COMMAND_FLAG == True:
                Log.info("Executing command on localhost...")
                os.system(command)

if __name__ == "__main__":
    main()
